# -*- coding: utf-8 -*-

"""详解见书: G = (V, E, g)"""

class Dijkstra():
    def __init__(self) -> None:
        pass

    def shortest_length(self, vertex, edge, g, u0):
        """输入：
        vertex = {0,1,2,3,4,5,6} 共7个按编号的顶点集，因为边和权都用二维矩阵表示，所以顶点可以直接用集合的基来表示，即 7
        edge   = [[],[],[]] 7*7 布尔矩阵，表示顶点 {u,v} 之间是否有边邻接
        g      = [[],[],[]] 7*7 数字矩阵，表示边 {u,v} 的权
        u0 表示计算顶点 u0 到其余所有顶点的最短路径，例如 3，表示从 3 号顶点出发
        --------------------------------------------------------------------
        过程中：
        left = {1,5,6} 还未计算出最短路径的顶点集
        leftC = {0,2,4} 已经计算出最短路径的顶点集, 与 left 互为补集
        shortest = [1.5, inf, 2.3, 0, 3.6, inf, inf] 从3号顶点到其他顶点的最短路径值
        previous = [3, None, 3, None, 3, None, None] 最断路径上的前导顶点是哪个，方便路径追溯"""

        left, leftC = set(range(vertex)), set()
        left.discard(u0)

        shortest = [float('inf') for _ in range(vertex)]
        shortest[u0] = 0
        previous = [None for _ in range(vertex)]

        while left:
            u1, v1, min_length = None, None, float('inf')
            for u in leftC:
                for v in range(vertex):
                    if edge[u][v] and v in left:
                        length = shortest[u] + g[u][v]
                        if length < min_length:
                            u1, v1, min_length = u, v, length
            for v in range(vertex):
                if edge[u0][v] and v in left:
                    if g[u0][v] < min_length:
                        u1, v1, min_length = u0, v, g[u0][v]

            shortest[v1] = min_length
            previous[v1] = u1
            left.discard(v1)
            leftC.add(v1)

        return shortest, previous

    def find_path(self, u, previous):
        """根据前置顶点找出最短的路径"""
        path = [u]
        def do(u):
            u1 = previous[u]
            if u1 is not None:
                path.append(u1)
                do(u1)
            else:
                return
        do(u)
        path.reverse()
        return path


if __name__ == '__main__':
    dijkstra = Dijkstra()
    vertex = 4
    edge = [[0,1,1,1],[1,0,1,0],[1,1,0,1],[1,0,1,0]]
    g = [[0,6,2,2],[6,0,2,float('inf')],[2,2,0,3],[2,float('inf'),3,0]]
    u0 = 1

    shortest, previous = dijkstra.shortest_length(vertex, edge, g, u0)

    print(shortest)
    print(previous)

    for u in range(vertex):
        if u != u0:
            path = dijkstra.find_path(u, previous)
            print('u:', u, ', path:', path)
